from typing import Dict, List, Optional

import aiohttp


class Query:
    """A class to interact with the AtomGit paper search API.

    This class provides methods to search and retrieve papers using different criteria
    such as content, paper ID, and title.
    """

    def __init__(self, *args, **kwargs):
        """Initialize the Query instance.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.
        """
        self.base_url = "http://180.184.65.98:38880/atomgit"

    async def query_by_content(
        self, query: str, top_k: int = 5, similarity_threshold: Optional[float] = None
    ) -> List[Dict]:
        """Search papers based on content query.

        Args:
            query (str): The search query string.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.
            similarity_threshold (float, optional): Minimum similarity score to return. Defaults to None.
        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/search_papers"
        params = {"query": query, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    results = await response.json()
                    if similarity_threshold:
                        results = [
                            result
                            for result in results
                            if result["distance"] <= similarity_threshold
                        ]
                    return results
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_by_id(self, paper_id: str, top_k: int = 50) -> List[Dict]:
        """Search papers based on paper ID.

        Args:
            paper_id (str): The ID of the paper to search for.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_paper_id"
        params = {"paper_id": paper_id, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_by_title(self, title: str, top_k: int = 5) -> List[Dict]:
        """Search papers based on paper title.

        Args:
            title (str): The title of the paper to search for.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_title"
        params = {"title": title, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_by_title_contain(
        self, substring: str, top_k: int = 5
    ) -> List[Dict]:
        """Search papers based on substring match in titles.

        Args:
            substring (str): The substring to search for in paper titles.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_title_contain"
        params = {"title": substring, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_by_chunk_contain(
        self, substring: str, top_k: int = 5
    ) -> List[Dict]:
        """Search papers based on substring match in content chunks.

        Args:
            substring (str): The substring to search for in paper content chunks.
            top_k (int, optional): Maximum number of results to return. Defaults to 5.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_chunk_contain"
        params = {"chunk": substring, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_by_keyword(
        self, keyword: str, top_k: int = 5
    ) -> List[Dict]:
        """Search papers based on keyword.

        Args:
            keyword (str): The keyword to search for in paper content chunks.

        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/query_by_keyword"
        params = {"keyword": keyword, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []


    async def query_by_title_like(
        self, title: str, top_k: int = 5
    ) -> List[Dict]:
        """通过相似标题查询论文，并返回相关的论文片段。

        Args:
            title (字符串, 必填): 要查询的标题关键词。
            top_k (整数, 可选, 默认值: 1): 查找最相似的标题数量 (用于初步筛选)。

        Returns:
            List[Dict]:返回值: 嵌套的论文片段信息列表。 
            外层列表对应每个相似标题，内层列表是与该相似标题相关的论文片段信息列表 (结构同 /search_papers)。

        """
        query_url = f"{self.base_url}/query_by_title_like"
        params = {"title": title, "top_k": top_k}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []


    async def query_keywords_metadata(
        self
    ) -> List[Dict]:
        """查看关键词统计
        Returns:
            List[Dict]: A list of dictionaries containing paper information.
                       Returns empty list if the request fails.
        """
        query_url = f"{self.base_url}/keywords_metadata"
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return []

    async def query_whole_text_by_id(self, paper_id: str) -> Optional[str]:
        """通过论文ID查询论文全文。

        Args:
            paper_id (str): 要查询的论文ID。

        Returns:
            Optional[str]: 论文完整文本，如果未找到则返回None。
                         如果请求失败则返回None。
        """
        query_url = f"{self.base_url}/query_whole_text_by_id"
        params = {"paper_id": paper_id}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.text()
        except Exception as e:
            print(f"Error making request: {e}")
            return None

    async def query_whole_text_by_title(self, title: str) -> Optional[str]:
        """通过论文标题查询论文全文。

        Args:
            title (str): 要查询的论文标题。

        Returns:
            Optional[str]: 论文完整文本，如果未找到则返回None。
                         如果请求失败则返回None。
        """
        query_url = f"{self.base_url}/query_whole_text_by_title"
        params = {"title": title}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.text()
        except Exception as e:
            print(f"Error making request: {e}")
            return None

    async def query_keywords_by_id(self, paper_id: str) -> Optional[List[str]]:
        """通过论文ID查询论文关键词。

        Args:
            paper_id (str): 要查询的论文ID。

        Returns:
            Optional[List[str]]: 关键词列表，如果未找到则返回None。
                               如果请求失败则返回None。
        """
        query_url = f"{self.base_url}/query_keywords_by_id"
        params = {"paper_id": paper_id}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return None

    async def query_keywords_by_title(self, title: str) -> Optional[List[str]]:
        """通过论文标题查询论文关键词。

        Args:
            title (str): 要查询的论文标题。

        Returns:
            Optional[List[str]]: 关键词列表，如果未找到则返回None。
                               如果请求失败则返回None。
        """
        query_url = f"{self.base_url}/query_keywords_by_title"
        params = {"title": title}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(query_url, params=params) as response:
                    response.raise_for_status()
                    return await response.json()
        except Exception as e:
            print(f"Error making request: {e}")
            return None



if __name__ == "__main__":
    import asyncio

    async def main():
        query = Query()
        query = Query()
        print(await query.query_by_keyword("deep learning"))
        print(await query.query_keywords_metadata())
        print(await query.query_by_title_like("Imbedding Deep Neural Networks"))
        print(await query.query_by_title_contain("deep learning"))
        print(await query.query_by_chunk_contain("deep learning"))
        print(await query.query_by_id("651b7dbc3fda6d7f06304579"))
        print(await query.query_by_title("SeA: Semantic Adversarial Augmentation for Last Layer Features from Unsupervised Representation Learning"))
        print(await query.query_by_content("deep learning"))
        print(await query.query_by_title_like("Imbedding Deep Neural Networks"))
        print(await query.query_by_title_contain("deep learning"))
        print(await query.query_by_chunk_contain("deep learning"))
        print(await query.query_by_id("651b7dbc3fda6d7f06304579"))
        print(await query.query_by_title("SeA: Semantic Adversarial Augmentation for Last Layer Features from Unsupervised Representation Learning"))
        print(await query.query_by_content("deep learning"))
    asyncio.run(main())
